import os
try:
    from .cvio import cvio
except:
    import sys
    sys.path.append(os.getcwd())
    from function.cvio import cvio


def getMergeDic(src):
    pairs = {}
    with open(src, encoding='utf-8') as fp:
        for f in fp.readlines():
            f = f.replace('\n', '').replace('\t',',').replace('，',',')#.replace(' ',',')
            s = f.split(',')
            pairs[s[0]] = s[-1].split('|')[-1]
    return pairs

def rename_maskset(src, txt=None, spec=None, recursive=False):
    if txt in (None, '') and spec in (None, ''):
        print('结束执行，无有效输入')
        return
    imgannlist = cvio.load_img_ann_list(src, recursive=recursive, silent=False)
    if not len(imgannlist):
        imgannlist = cvio.load_ext_list(src, recursive=recursive, silent=False)
        imgannlist = [[ann, ann] for ann in imgannlist]
    if  os.path.exists(txt):
        ldic = getMergeDic(txt)
    else:
        ldic = {}

    n = len(imgannlist)
    for i, (img, ann) in enumerate(imgannlist, 1):
        print('[%d/%d]' % (i, n), os.path.basename(img))
        annInfo = cvio.load_ann(ann)
        F = False
        for shape in annInfo['shapes']:
            label = shape['label']
            if label in ldic:
                shape['label'] = ldic[label]
                F = True
            elif spec:
                shape['label'] = spec
                F = True
        if F:
            cvio.write_ann(annInfo, ann)

if __name__ == '__main__':
    src = r'J:\hjz\train_cj\20221021_finished' # 
    txt = r'J:\hjz\train_cj\20221021_finished\map.txt'
    rename_maskset(src, txt, recursive=True)
    # pyinstaller -F -i m13.ico main.py
